#!/bin/bash
set -euo pipefail

# Change to root directory of repo
DIRNAME=$(cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
cd "${DIRNAME}/.."

# Get current git head & stash unsaved changes
GIT_HEAD=$(git rev-parse HEAD)
GIT_BRANCH=$(git branch --show-current)
if [ -z "${GIT_BRANCH}" ]; then
    GIT_BRANCH="main"
fi
n_stash=$(git stash list | wc -l)
git stash

# Set up exit handler to restore git state & delete temp branches
function exit_handler {
    git reset --hard && git checkout "${GIT_BRANCH}" --
    if (( n_stash < $(git stash list | wc -l) )); then
        git stash pop
    fi
    for version in $(git tag --list 'v[0-9]*'); do
        branch="${version}_local_docs_only"
        if git show-ref --verify --quiet "refs/heads/$branch"; then
            git branch -D "$branch"
        fi
    done
}
trap exit_handler EXIT

# Set up virtual environment & install Sphinx requirements.
pip3 install --upgrade pip setuptools wheel virtualenv
if [ ! -d venv ]; then
  rm -f venv
  virtualenv venv
fi
source venv/bin/activate
pip3 install -r "${DIRNAME}/requirements.txt"

# Get old Merlion docs from gh-pages branch. Only keep version-tagged ones, and update the version matrix as needed.
git checkout gh-pages && git pull && git checkout --force "${GIT_HEAD}"
sphinx-build -M clean "${DIRNAME}/source" "${DIRNAME}/build"
mkdir -p "${DIRNAME}/build" "${DIRNAME}/build/html"
git --work-tree "${DIRNAME}/build/html" checkout gh-pages . && git reset --hard
python3 "${DIRNAME}/process_old_docs.py"

# Install all released versions of Merlion/ts_datasets _not_ on gh-pages and use them to build the appropriate API docs.
# Uninstall after we're done with each one.
versions=("latest")
for v in $(git tag --list 'v[0-9]*'); do
  versions+=("$v")
done
docs_files=("${DIRNAME}/source/*.rst" "examples")
checkout_files=("${docs_files[@]}" "merlion" "ts_datasets" "setup.py" "MANIFEST.in")
for version in "${versions[@]}"; do
    if [[ ! -d "${DIRNAME}/build/html/${version}" ]]; then
        rm -rf "${DIRNAME}/build/html/${version}"
        if [[ ${version} != "latest" ]]; then
            git checkout -b "${version}_local_docs_only"
            for f in $(git diff --name-only --diff-filter=AR "tags/${version}" "${DIRNAME}/source/*.rst" "examples"); do
                git rm "$f"
            done
            git checkout "tags/${version}" -- "${checkout_files[@]}"
        fi
        export current_version=${version}
        pip3 install ".[all]"
        pip3 install ts_datasets/
        sphinx-build -b html "${DIRNAME}/source" "${DIRNAME}/build/html/${current_version}" -W --keep-going
        rm -rf "${DIRNAME}/build/html/${current_version}/.doctrees"
        pip3 uninstall -y salesforce-merlion ts_datasets
        git reset --hard && git checkout --force "${GIT_HEAD}"
    fi
done

# Determine the latest stable version if there is one
if (( ${#versions[@]} > 1 )); then
  stable_hash=$(git rev-list --tags --max-count=1)
  stable_version=$(git describe --tags "$stable_hash")
  export stable_version
else
  export stable_version="latest"
fi

# Create dummy HTML's for the stable version in the base directory
while read -r filename; do
    filename=$(echo "$filename" | sed "s/\.\///")
    n_sub=$(echo "$filename" | (grep -o "/" || true) | wc -l)
    prefix=""
    for (( i=0; i<n_sub; i++ )); do
        prefix+="../"
    done
    url="${prefix}${stable_version}/$filename"
    mkdir -p "${DIRNAME}/build/html/$(dirname "$filename")"
    cat > "${DIRNAME}/build/html/$filename" <<EOF
<!DOCTYPE html>
<html>
   <head>
      <title>Merlion Documentation</title>
      <meta http-equiv = "refresh" content="0; url='$url'" />
   </head>
   <body>
      <p>Please wait while you're redirected to our <a href="$url">documentation</a>.</p>
   </body>
</html>
EOF
done < <(cd "${DIRNAME}/build/html/$stable_version" && find . -name "*.html")

# Add README
cat > "${DIRNAME}/build/html/README.md" <<EOF
# GitHub Pages

The contents of this branch are automatically generated by GitHub Actions. This branch is used by GitHub Pages to
populate the API documentation [here](https://opensource.salesforce.com/Merlion/index.html).

EOF
echo "Finished writing to build/html."
